from langchain import hub
from langchain.agents import AgentType, Tool, initialize_agent, AgentExecutor
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.agents.output_parsers import SelfAskOutputParser
from langchain.llms import OpenAI
from langchain.utilities import SerpAPIWrapper

llm = OpenAI(temperature=0)
search = SerpAPIWrapper()
tools = [
    Tool(
        name="Intermediate Answer",
        func=search.run,
        description="useful for when you need to ask with search",
    )
]
prompt = hub.pull("hwchase17/self-ask-with-search")
llm_with_stop = llm.bind(stop=["\nIntermediate answer:"])
agent = (
    {
        "input": lambda x: x["input"],
        # Use some custom observation_prefix/llm_prefix for formatting
        "agent_scratchpad": lambda x: format_log_to_str(
            x["intermediate_steps"],
            observation_prefix="\nIntermediate answer: ",
            llm_prefix="",
        ),
    }
    | prompt
    | llm_with_stop
    | SelfAskOutputParser()
)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

agent_executor.invoke(
    {"input": "What is the hometown of the reigning men's U.S. Open champion?"}
)

#使用现成的代理

self_ask_with_search = initialize_agent(
    tools, llm, agent=AgentType.SELF_ASK_WITH_SEARCH, verbose=True
)
self_ask_with_search.run(
    "What is the hometown of the reigning men's U.S. Open champion?"
)